{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "narrative-sailing",
   "metadata": {
    "papermill": {
     "duration": 0.040767,
     "end_time": "2021-03-24T14:13:51.082088",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.041321",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Inference on Predictive and Causal Effects in High-Dimensional Nonlinear Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ready-appearance",
   "metadata": {
    "papermill": {
     "duration": 0.036878,
     "end_time": "2021-03-24T14:13:51.155414",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.118536",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Impact of 401(k) on  Financial Wealth\n",
    "\n",
    "As a practical illustration of the methods developed in this lecture, we consider estimation of the effect of 401(k) eligibility and participation \n",
    "on accumulated assets. 401(k) plans are pension accounts sponsored by employers. The key problem in determining the effect of participation in 401(k) plans on accumulated assets is saver heterogeneity coupled with the fact that the decision to enroll in a 401(k) is non-random. It is generally recognized that some people have a higher preference for saving than others. It also seems likely that those individuals with high unobserved preference for saving would be most likely to choose to participate in tax-advantaged retirement savings plans and would tend to have otherwise high amounts of accumulated assets. The presence of unobserved savings preferences with these properties then implies that conventional estimates that do not account for saver heterogeneity and endogeneity of participation will be biased upward, tending to overstate the savings effects of 401(k) participation.\n",
    "\n",
    "One can argue that eligibility for enrolling in a 401(k) plan in this data can be taken as exogenous after conditioning on a few observables of which the most important for their argument is income. The basic idea is that, at least around the time 401(k)’s initially became available, people were unlikely to be basing their employment decisions on whether an employer offered a 401(k) but would instead focus on income and other aspects of the job. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "divine-phoenix",
   "metadata": {
    "papermill": {
     "duration": 0.036497,
     "end_time": "2021-03-24T14:13:51.229370",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.192873",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Data\n",
    "\n",
    "The data set can be loaded from the `hdm` package for R by typing\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "concerned-continent",
   "metadata": {
    "papermill": {
     "duration": 0.481645,
     "end_time": "2021-03-24T14:13:51.748472",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.266827",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "library(hdm) \n",
    "library(ggplot2)\n",
    "data(pension)\n",
    "data <- pension \n",
    "dim(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "horizontal-sheffield",
   "metadata": {
    "papermill": {
     "duration": 0.036873,
     "end_time": "2021-03-24T14:13:51.825239",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.788366",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "See the \"Details\" section on the description of the data set, which can be accessed by \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "absolute-intersection",
   "metadata": {
    "papermill": {
     "duration": 0.311328,
     "end_time": "2021-03-24T14:13:52.174594",
     "exception": false,
     "start_time": "2021-03-24T14:13:51.863266",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "help(pension)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "looking-invention",
   "metadata": {
    "papermill": {
     "duration": 0.037509,
     "end_time": "2021-03-24T14:13:52.250583",
     "exception": false,
     "start_time": "2021-03-24T14:13:52.213074",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The data consist of 9,915 observations at the household level drawn from the 1991 Survey of Income and Program Participation (SIPP).  All the variables are referred to 1990. We use net financial assets (*net\\_tfa*) as the outcome variable, $Y$,  in our analysis. The net financial assets are computed as the sum of IRA balances, 401(k) balances, checking accounts, saving bonds, other interest-earning accounts, other interest-earning assets, stocks, and mutual funds less non mortgage debts. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "received-nutrition",
   "metadata": {
    "papermill": {
     "duration": 0.038843,
     "end_time": "2021-03-24T14:13:52.327055",
     "exception": false,
     "start_time": "2021-03-24T14:13:52.288212",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Among the $9915$ individuals, $3682$ are eligible to participate in the program. The variable *e401* indicates eligibility and *p401* indicates participation, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "continuing-composite",
   "metadata": {
    "papermill": {
     "duration": 0.991238,
     "end_time": "2021-03-24T14:13:53.355568",
     "exception": false,
     "start_time": "2021-03-24T14:13:52.364330",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "hist_e401 = ggplot(data, aes(x = e401, fill = factor(e401))) + geom_bar()\n",
    "hist_e401"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "material-sending",
   "metadata": {
    "papermill": {
     "duration": 0.039406,
     "end_time": "2021-03-24T14:13:53.435306",
     "exception": false,
     "start_time": "2021-03-24T14:13:53.395900",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Eligibility is highly associated with financial wealth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "legislative-patrol",
   "metadata": {
    "papermill": {
     "duration": 0.688917,
     "end_time": "2021-03-24T14:13:54.163799",
     "exception": false,
     "start_time": "2021-03-24T14:13:53.474882",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dens_net_tfa = ggplot(data, aes(x = net_tfa, color = factor(e401), fill = factor(e401)) ) + \n",
    "                    geom_density() + xlim(c(-20000, 150000)) + \n",
    "                    facet_wrap(.~e401) \n",
    "                    \n",
    "dens_net_tfa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "awful-antigua",
   "metadata": {
    "papermill": {
     "duration": 0.040645,
     "end_time": "2021-03-24T14:13:54.245822",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.205177",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The unconditional APE of e401 is about $19559$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "interesting-jones",
   "metadata": {
    "papermill": {
     "duration": 0.068689,
     "end_time": "2021-03-24T14:13:54.354643",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.285954",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "e1 <- data[data$e401==1,]\n",
    "e0 <- data[data$e401==0,]\n",
    "round(mean(e1$net_tfa)-mean(e0$net_tfa),0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cross-priority",
   "metadata": {
    "papermill": {
     "duration": 0.041187,
     "end_time": "2021-03-24T14:13:54.437695",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.396508",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Among the $3682$ individuals that  are eligible, $2594$ decided to participate in the program. The unconditional APE of p401 is about $27372$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "organized-radius",
   "metadata": {
    "papermill": {
     "duration": 0.078062,
     "end_time": "2021-03-24T14:13:54.557173",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.479111",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "p1 <- data[data$p401==1,]\n",
    "p0 <- data[data$p401==0,]\n",
    "round(mean(p1$net_tfa)-mean(p0$net_tfa),0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "suitable-vulnerability",
   "metadata": {
    "papermill": {
     "duration": 0.041678,
     "end_time": "2021-03-24T14:13:54.640532",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.598854",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "As discussed, these estimates are biased since they do not account for saver heterogeneity and endogeneity of participation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "balanced-damages",
   "metadata": {
    "papermill": {
     "duration": 0.041283,
     "end_time": "2021-03-24T14:13:54.723664",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.682381",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Double ML package"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "recent-southwest",
   "metadata": {
    "papermill": {
     "duration": 0.041034,
     "end_time": "2021-03-24T14:13:54.806271",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.765237",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We are interested in valid estimators of the average treatment effect of `e401` and `p401` on `net_tfa`. To get those estimators, we use the `DoubleML` package that internally builds on mlr3. You find additional information on the package on the package website https://docs.doubleml.org/ and the R documentation page https://docs.doubleml.org/r/stable/. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "popular-matrix",
   "metadata": {
    "papermill": {
     "duration": 170.599415,
     "end_time": "2021-03-24T14:16:45.447305",
     "exception": false,
     "start_time": "2021-03-24T14:13:54.847890",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# installing Double ML\n",
    "remotes::install_github(\"DoubleML/doubleml-for-r\",quiet=TRUE)\n",
    "\n",
    "\n",
    "# loading the packages\n",
    "library(DoubleML)\n",
    "library(mlr3learners)\n",
    "library(mlr3)\n",
    "library(data.table)\n",
    "library(randomForest)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "empirical-basics",
   "metadata": {
    "papermill": {
     "duration": 0.045129,
     "end_time": "2021-03-24T14:16:45.537784",
     "exception": false,
     "start_time": "2021-03-24T14:16:45.492655",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "As mentioned, in the tutorial we use the meta package `mlr3` to generate predictions with machine learning methods. A comprehensive introduction and description of the `mlr3` package is provided in the [mlr3book](https://mlr3book.mlr-org.com/). A list of all learners that you can use in `mlr3` can be found [here](https://mlr3extralearners.mlr-org.com/articles/learners/list_learners.html). The entry in the columns *mlr3 Package* and *Packages* indicate which packages must be installed/loaded in your R session. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "other-leader",
   "metadata": {
    "papermill": {
     "duration": 0.044769,
     "end_time": "2021-03-24T14:16:45.625738",
     "exception": false,
     "start_time": "2021-03-24T14:16:45.580969",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Estimating the ATE of 401(k) Eligibility on Net Financial Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "reliable-solid",
   "metadata": {
    "papermill": {
     "duration": 0.043232,
     "end_time": "2021-03-24T14:16:45.711972",
     "exception": false,
     "start_time": "2021-03-24T14:16:45.668740",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We first look at the treatment effect of e401 on net total financial assets. We give estimates of the ATE and ATT that corresponds to the linear model\n",
    "\n",
    "\\begin{equation*}\n",
    "Y = D \\alpha + f(X)'\\beta+ \\epsilon,\n",
    "\\end{equation*}\n",
    "\n",
    "where $f(X)$ includes indicators of marital status, two-earner status, defined benefit pension status, IRA participation status, and home ownership status, and  orthogonal polynomials of degrees 2, 4, 6 and 8 in family size, education, age and  income, respectively. The dimensions of $f(X)$ is 25. \n",
    "\n",
    "In the first step, we report estimates of the average treatment effect (ATE) of 401(k) eligibility on net financial assets both in the partially linear regression (PLR) model and in the interactive regression model (IRM) allowing for heterogeneous treatment effects. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "directed-discovery",
   "metadata": {
    "papermill": {
     "duration": 0.328447,
     "end_time": "2021-03-24T14:16:46.084312",
     "exception": false,
     "start_time": "2021-03-24T14:16:45.755865",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Constructing the data (as DoubleMLData)\n",
    "formula_flex = \"net_tfa ~ e401 + poly(age, 6, raw=TRUE) + poly(inc, 8, raw=TRUE) + poly(educ, 4, raw=TRUE) + poly(fsize, 2, raw=TRUE) + marr + twoearn + db + pira + hown\"\n",
    "model_flex = as.data.table(model.frame(formula_flex, pension))\n",
    "x_cols = colnames(model_flex)[-c(1,2)]\n",
    "data_ml = DoubleMLData$new(model_flex, y_col = \"net_tfa\", d_cols = \"e401\", x_cols=x_cols)\n",
    "\n",
    "\n",
    "p <- dim(model_flex)[2]-2\n",
    "p\n",
    "\n",
    "# complex model with two-way interactions\n",
    "#data_interactions = fetch_401k(polynomial_features = TRUE, instrument = FALSE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "weekly-advancement",
   "metadata": {
    "papermill": {
     "duration": 0.047299,
     "end_time": "2021-03-24T14:16:46.176537",
     "exception": false,
     "start_time": "2021-03-24T14:16:46.129238",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Partially Linear Regression Models (PLR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "united-massachusetts",
   "metadata": {
    "papermill": {
     "duration": 0.044606,
     "end_time": "2021-03-24T14:16:46.266797",
     "exception": false,
     "start_time": "2021-03-24T14:16:46.222191",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We start using lasso to estimate the function $g_0$ and $m_0$ in the following PLR model:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interesting-symphony",
   "metadata": {
    "papermill": {
     "duration": 0.044208,
     "end_time": "2021-03-24T14:16:46.355767",
     "exception": false,
     "start_time": "2021-03-24T14:16:46.311559",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "\\begin{eqnarray}\n",
    " &  Y = D\\theta_0 + g_0(X) + \\zeta,  &  E[\\zeta \\mid D,X]= 0,\\\\\n",
    " & D = m_0(X) +  V,   &  E[V \\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "entire-killer",
   "metadata": {
    "papermill": {
     "duration": 22.997053,
     "end_time": "2021-03-24T14:17:09.397895",
     "exception": false,
     "start_time": "2021-03-24T14:16:46.400842",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Estimating the PLR\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "set.seed(123)\n",
    "lasso <- lrn(\"regr.cv_glmnet\",nfolds = 5, s = \"lambda.min\")\n",
    "lasso_class <- lrn(\"classif.cv_glmnet\", nfolds = 5, s = \"lambda.min\")\n",
    "\n",
    "dml_plr <- DoubleMLPLR$new(data_ml, ml_g = lasso, ml_m = lasso_class, n_folds=3)\n",
    "dml_plr$fit(store_predictions=TRUE)\n",
    "dml_plr$summary()\n",
    "lasso_plr <- dml_plr$coef\n",
    "lasso_std_plr <- dml_plr$se"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prerequisite-browser",
   "metadata": {
    "papermill": {
     "duration": 0.04484,
     "end_time": "2021-03-24T14:17:09.488491",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.443651",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let us check the predictive performance of this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "middle-spokesman",
   "metadata": {
    "papermill": {
     "duration": 0.069669,
     "end_time": "2021-03-24T14:17:09.602972",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.533303",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dml_plr$params_names()\n",
    "g_hat <- as.matrix(dml_plr$predictions$ml_g) # predictions of g_o\n",
    "m_hat <- as.matrix(dml_plr$predictions$ml_m) # predictions of m_o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wrong-letters",
   "metadata": {
    "papermill": {
     "duration": 0.075135,
     "end_time": "2021-03-24T14:17:09.726237",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.651102",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# cross-fitted RMSE: outcome\n",
    "y <- as.matrix(pension$net_tfa) # true observations\n",
    "theta <- as.numeric(dml_plr$coef) # estimated regression coefficient\n",
    "d <- as.matrix(pension$e401) \n",
    "predictions_y <- as.matrix(d*theta)+g_hat # predictions for y\n",
    "lasso_y_rmse <- sqrt(mean((y-predictions_y)^2)) \n",
    "lasso_y_rmse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "medieval-friend",
   "metadata": {
    "papermill": {
     "duration": 0.075492,
     "end_time": "2021-03-24T14:17:09.848578",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.773086",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# cross-fitted RMSE: treatment\n",
    "d <- as.matrix(pension$e401) \n",
    "lasso_d_rmse <- sqrt(mean((d-m_hat)^2)) \n",
    "lasso_d_rmse\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fluid-gregory",
   "metadata": {
    "papermill": {
     "duration": 0.048205,
     "end_time": "2021-03-24T14:17:09.946894",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.898689",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Then, we repeat this procedure for various machine learning methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "naughty-moscow",
   "metadata": {
    "papermill": {
     "duration": 62.65345,
     "end_time": "2021-03-24T14:18:12.648324",
     "exception": false,
     "start_time": "2021-03-24T14:17:09.994874",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Random Forest\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "randomForest <- lrn(\"regr.ranger\")\n",
    "randomForest_class <- lrn(\"classif.ranger\")\n",
    "\n",
    "dml_plr <- DoubleMLPLR$new(data_ml, ml_g = randomForest, ml_m = randomForest_class, n_folds=3)\n",
    "dml_plr$fit(store_predictions=TRUE) # set store_predictions=TRUE to evaluate the model\n",
    "dml_plr$summary()\n",
    "forest_plr <- dml_plr$coef\n",
    "forest_std_plr <- dml_plr$se"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "legislative-balance",
   "metadata": {
    "papermill": {
     "duration": 0.04819,
     "end_time": "2021-03-24T14:18:12.746400",
     "exception": false,
     "start_time": "2021-03-24T14:18:12.698210",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can compare the accuracy of this model to the model that has been estimated with lasso."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "described-expansion",
   "metadata": {
    "papermill": {
     "duration": 0.132718,
     "end_time": "2021-03-24T14:18:12.927195",
     "exception": false,
     "start_time": "2021-03-24T14:18:12.794477",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Evaluation predictions\n",
    "g_hat <- as.matrix(dml_plr$predictions$ml_g) # predictions of g_o\n",
    "m_hat <- as.matrix(dml_plr$predictions$ml_m) # predictions of m_o\n",
    "theta <- as.numeric(dml_plr$coef) # estimated regression coefficient\n",
    "predictions_y <- as.matrix(d*theta)+g_hat # predictions for y\n",
    "forest_y_rmse <- sqrt(mean((y-predictions_y)^2)) \n",
    "forest_y_rmse\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "forest_d_rmse <- sqrt(mean((d-m_hat)^2)) \n",
    "forest_d_rmse\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aquatic-toner",
   "metadata": {
    "papermill": {
     "duration": 1.567333,
     "end_time": "2021-03-24T14:18:14.545412",
     "exception": false,
     "start_time": "2021-03-24T14:18:12.978079",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Trees\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "\n",
    "trees <- lrn(\"regr.rpart\")\n",
    "trees_class <- lrn(\"classif.rpart\")\n",
    "\n",
    "dml_plr <- DoubleMLPLR$new(data_ml, ml_g = trees, ml_m = trees_class, n_folds=3)\n",
    "dml_plr$fit(store_predictions=TRUE)\n",
    "dml_plr$summary()\n",
    "tree_plr <- dml_plr$coef\n",
    "tree_std_plr <- dml_plr$se\n",
    "\n",
    "# Evaluation predictions\n",
    "g_hat <- as.matrix(dml_plr$predictions$ml_g) # predictions of g_o\n",
    "m_hat <- as.matrix(dml_plr$predictions$ml_m) # predictions of m_o\n",
    "theta <- as.numeric(dml_plr$coef) # estimated regression coefficient\n",
    "predictions_y <- as.matrix(d*theta)+g_hat # predictions for y\n",
    "tree_y_rmse <- sqrt(mean((y-predictions_y)^2)) \n",
    "tree_y_rmse\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "tree_d_rmse <- sqrt(mean((d-m_hat)^2)) \n",
    "tree_d_rmse\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collectible-engineering",
   "metadata": {
    "papermill": {
     "duration": 1.005937,
     "end_time": "2021-03-24T14:18:15.603603",
     "exception": false,
     "start_time": "2021-03-24T14:18:14.597666",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Boosting\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "boost<- lrn(\"regr.xgboost\",objective=\"reg:squarederror\")\n",
    "boost_class <- lrn(\"classif.xgboost\",objective = \"binary:logistic\",eval_metric =\"logloss\")\n",
    "\n",
    "dml_plr <- DoubleMLPLR$new(data_ml, ml_g = boost, ml_m = boost_class, n_folds=3)\n",
    "dml_plr$fit(store_predictions=TRUE)\n",
    "dml_plr$summary()\n",
    "boost_plr <- dml_plr$coef\n",
    "boost_std_plr <- dml_plr$se\n",
    "\n",
    "# Evaluation predictions\n",
    "g_hat <- as.matrix(dml_plr$predictions$ml_g) # predictions of g_o\n",
    "m_hat <- as.matrix(dml_plr$predictions$ml_m) # predictions of m_o\n",
    "theta <- as.numeric(dml_plr$coef) # estimated regression coefficient\n",
    "predictions_y <- as.matrix(d*theta)+g_hat # predictions for y\n",
    "boost_y_rmse <- sqrt(mean((y-predictions_y)^2)) \n",
    "boost_y_rmse\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "boost_d_rmse <- sqrt(mean((d-m_hat)^2)) \n",
    "boost_d_rmse\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "laden-netherlands",
   "metadata": {
    "papermill": {
     "duration": 0.057943,
     "end_time": "2021-03-24T14:18:15.717190",
     "exception": false,
     "start_time": "2021-03-24T14:18:15.659247",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's sum up the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "seven-hawaii",
   "metadata": {
    "papermill": {
     "duration": 0.158103,
     "end_time": "2021-03-24T14:18:15.930645",
     "exception": false,
     "start_time": "2021-03-24T14:18:15.772542",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "library(xtable)\n",
    "table <- matrix(0, 4, 4)\n",
    "table[1,1:4]   <- c(lasso_plr,forest_plr,tree_plr,boost_plr)\n",
    "table[2,1:4]   <- c(lasso_std_plr,forest_std_plr,tree_std_plr,boost_std_plr)\n",
    "table[3,1:4]   <- c(lasso_y_rmse,forest_y_rmse,tree_y_rmse,boost_y_rmse)\n",
    "table[4,1:4]   <- c(lasso_d_rmse,forest_d_rmse,tree_d_rmse,boost_d_rmse)\n",
    "rownames(table) <- c(\"Estimate\",\"Std.Error\",\"RMSE Y\",\"RMSE D\")\n",
    "colnames(table) <- c(\"Lasso\",\"Random Forest\",\"Trees\",\"Boosting\")\n",
    "tab<- xtable(table, digits = 2)\n",
    "tab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "identical-smith",
   "metadata": {
    "papermill": {
     "duration": 0.056641,
     "end_time": "2021-03-24T14:18:16.044675",
     "exception": false,
     "start_time": "2021-03-24T14:18:15.988034",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The best model with lowest RMSE in both equation is the PLR model estimated via lasso. It gives the following estimate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "occupied-wrong",
   "metadata": {
    "papermill": {
     "duration": 0.082029,
     "end_time": "2021-03-24T14:18:16.183273",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.101244",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "lasso_plr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "comprehensive-graphics",
   "metadata": {
    "papermill": {
     "duration": 0.056567,
     "end_time": "2021-03-24T14:18:16.296870",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.240303",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Interactive Regression Model (IRM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "square-craps",
   "metadata": {
    "papermill": {
     "duration": 0.057315,
     "end_time": "2021-03-24T14:18:16.410105",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.352790",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Next, we consider estimation of average treatment effects when treatment effects are fully heterogeneous:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "large-welcome",
   "metadata": {
    "papermill": {
     "duration": 0.05644,
     "end_time": "2021-03-24T14:18:16.525113",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.468673",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    " \\begin{eqnarray}\\label{eq: HetPL1}\n",
    " & Y  = g_0(D, X) + U,  &  \\quad E[U \\mid X, D]= 0,\\\\\n",
    "  & D  = m_0(X) + V,  & \\quad  E[V\\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "perfect-reliance",
   "metadata": {
    "papermill": {
     "duration": 0.056073,
     "end_time": "2021-03-24T14:18:16.637385",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.581312",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "To reduce the disproportionate impact of extreme propensity score weights in the interactive model\n",
    "we trim the propensity scores which are close to the bounds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "maritime-cleveland",
   "metadata": {
    "papermill": {
     "duration": 18.61409,
     "end_time": "2021-03-24T14:18:35.308502",
     "exception": false,
     "start_time": "2021-03-24T14:18:16.694412",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "dml_irm = DoubleMLIRM$new(data_ml, ml_g = lasso, \n",
    "                          ml_m = lasso_class, \n",
    "                          trimming_threshold = 0.01, n_folds=3)\n",
    "dml_irm$fit(store_predictions=TRUE)\n",
    "dml_irm$summary()\n",
    "lasso_irm <- dml_irm$coef\n",
    "lasso_std_irm <- dml_irm$se\n",
    "\n",
    "\n",
    "# predictions\n",
    "dml_irm$params_names()\n",
    "g0_hat <- as.matrix(dml_irm$predictions$ml_g0) # predictions of g_0(D=0, X)\n",
    "g1_hat <- as.matrix(dml_irm$predictions$ml_g1) # predictions of g_0(D=1, X)\n",
    "g_hat <- d*g1_hat+(1-d)*g0_hat # predictions of g_0\n",
    "m_hat <- as.matrix(dml_irm$predictions$ml_m) # predictions of m_o\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "artistic-parallel",
   "metadata": {
    "papermill": {
     "duration": 0.100289,
     "end_time": "2021-03-24T14:18:35.466922",
     "exception": false,
     "start_time": "2021-03-24T14:18:35.366633",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# cross-fitted RMSE: outcome\n",
    "y <- as.matrix(pension$net_tfa) # true observations\n",
    "d <- as.matrix(pension$e401) \n",
    "lasso_y_irm <- sqrt(mean((y-g_hat)^2)) \n",
    "lasso_y_irm\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "lasso_d_irm <- sqrt(mean((d-m_hat)^2)) \n",
    "lasso_d_irm\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hired-somerset",
   "metadata": {
    "papermill": {
     "duration": 59.894412,
     "end_time": "2021-03-24T14:19:35.420832",
     "exception": false,
     "start_time": "2021-03-24T14:18:35.526420",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "##### forest #####\n",
    "\n",
    "dml_irm = DoubleMLIRM$new(data_ml, ml_g = randomForest, \n",
    "                          ml_m = randomForest_class, \n",
    "                          trimming_threshold = 0.01, n_folds=3)\n",
    "dml_irm$fit(store_predictions=TRUE)\n",
    "dml_irm$summary()\n",
    "forest_irm <- dml_irm$coef\n",
    "forest_std_irm <- dml_plr$se\n",
    "\n",
    "# predictions\n",
    "g0_hat <- as.matrix(dml_irm$predictions$ml_g0) # predictions of g_0(D=0, X)\n",
    "g1_hat <- as.matrix(dml_irm$predictions$ml_g1) # predictions of g_0(D=1, X)\n",
    "g_hat <- d*g1_hat+(1-d)*g0_hat # predictions of g_0\n",
    "m_hat <- as.matrix(dml_irm$predictions$ml_m) # predictions of m_o\n",
    "\n",
    "# cross-fitted RMSE: outcome\n",
    "y <- as.matrix(pension$net_tfa) # true observations\n",
    "d <- as.matrix(pension$e401) \n",
    "forest_y_irm <- sqrt(mean((y-g_hat)^2)) \n",
    "forest_y_irm\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "forest_d_irm <- sqrt(mean((d-m_hat)^2)) \n",
    "forest_d_irm\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)\n",
    "\n",
    "##### trees #####\n",
    "\n",
    "dml_irm <- DoubleMLIRM$new(data_ml, ml_g = trees, ml_m = trees_class, \n",
    "                           trimming_threshold = 0.01, n_folds=3)\n",
    "dml_irm$fit(store_predictions=TRUE)\n",
    "dml_irm$summary()\n",
    "tree_irm <- dml_irm$coef\n",
    "tree_std_irm <- dml_irm$se\n",
    "\n",
    "# predictions\n",
    "g0_hat <- as.matrix(dml_irm$predictions$ml_g0) # predictions of g_0(D=0, X)\n",
    "g1_hat <- as.matrix(dml_irm$predictions$ml_g1) # predictions of g_0(D=1, X)\n",
    "g_hat <- d*g1_hat+(1-d)*g0_hat # predictions of g_0\n",
    "m_hat <- as.matrix(dml_irm$predictions$ml_m) # predictions of m_o\n",
    "\n",
    "# cross-fitted RMSE: outcome\n",
    "y <- as.matrix(pension$net_tfa) # true observations\n",
    "d <- as.matrix(pension$e401) \n",
    "tree_y_irm <- sqrt(mean((y-g_hat)^2)) \n",
    "tree_y_irm\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "tree_d_irm <- sqrt(mean((d-m_hat)^2)) \n",
    "tree_d_irm\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)\n",
    "\n",
    "\n",
    "##### boosting #####\n",
    "\n",
    "dml_irm <- DoubleMLIRM$new(data_ml, ml_g = boost, ml_m = boost_class,\n",
    "                           trimming_threshold = 0.01, n_folds=3)\n",
    "dml_irm$fit(store_predictions=TRUE)\n",
    "dml_irm$summary()\n",
    "boost_irm <- dml_irm$coef\n",
    "boost_std_irm <- dml_irm$se\n",
    "\n",
    "# predictions\n",
    "g0_hat <- as.matrix(dml_irm$predictions$ml_g0) # predictions of g_0(D=0, X)\n",
    "g1_hat <- as.matrix(dml_irm$predictions$ml_g1) # predictions of g_0(D=1, X)\n",
    "g_hat <- d*g1_hat+(1-d)*g0_hat # predictions of g_0\n",
    "m_hat <- as.matrix(dml_irm$predictions$ml_m) # predictions of m_o\n",
    "\n",
    "# cross-fitted RMSE: outcome\n",
    "y <- as.matrix(pension$net_tfa) # true observations\n",
    "d <- as.matrix(pension$e401) \n",
    "boost_y_irm <- sqrt(mean((y-g_hat)^2)) \n",
    "boost_y_irm\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "boost_d_irm <- sqrt(mean((d-m_hat)^2)) \n",
    "boost_d_irm\n",
    "\n",
    "# cross-fitted ce: treatment\n",
    "mean(ifelse(m_hat > 0.5, 1, 0) != d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "stopped-concrete",
   "metadata": {
    "papermill": {
     "duration": 0.109284,
     "end_time": "2021-03-24T14:19:35.601290",
     "exception": false,
     "start_time": "2021-03-24T14:19:35.492006",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "library(xtable)\n",
    "table <- matrix(0, 4, 4)\n",
    "table[1,1:4]   <- c(lasso_irm,forest_irm,tree_irm,boost_irm)\n",
    "table[2,1:4]   <- c(lasso_std_irm,forest_std_irm,tree_std_irm,boost_std_irm)\n",
    "table[3,1:4]   <- c(lasso_y_irm,forest_y_irm,tree_y_irm,boost_y_irm)\n",
    "table[4,1:4]   <- c(lasso_d_irm,forest_d_irm,tree_d_irm,boost_d_irm)\n",
    "rownames(table) <- c(\"Estimate\",\"Std.Error\",\"RMSE Y\",\"RMSE D\")\n",
    "colnames(table) <- c(\"Lasso\",\"Random Forest\",\"Trees\",\"Boosting\")\n",
    "tab<- xtable(table, digits = 2)\n",
    "tab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "surrounded-hospital",
   "metadata": {
    "papermill": {
     "duration": 0.067113,
     "end_time": "2021-03-24T14:19:35.736036",
     "exception": false,
     "start_time": "2021-03-24T14:19:35.668923",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Here, Random Forest gives the best prediction rule for $g_0$ and Lasso the best prediction rule for $m_0$, respectively. Let us fit the IRM model using the best ML method for each equation to get a final estimate for the treatment effect of eligibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "later-fraud",
   "metadata": {
    "papermill": {
     "duration": 45.616712,
     "end_time": "2021-03-24T14:20:21.422243",
     "exception": false,
     "start_time": "2021-03-24T14:19:35.805531",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "dml_irm = DoubleMLIRM$new(data_ml, ml_g = randomForest, \n",
    "                          ml_m = lasso_class, \n",
    "                          trimming_threshold = 0.01, n_folds=3)\n",
    "dml_irm$fit(store_predictions=TRUE)\n",
    "dml_irm$summary()\n",
    "best_irm <- dml_irm$coef\n",
    "best_std_irm <- dml_irm$se"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "discrete-karma",
   "metadata": {
    "papermill": {
     "duration": 0.066697,
     "end_time": "2021-03-24T14:20:21.555921",
     "exception": false,
     "start_time": "2021-03-24T14:20:21.489224",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "These estimates that flexibly account for confounding are\n",
    "substantially attenuated relative to the baseline estimate (*19559*) that does not account for confounding. They suggest much smaller causal effects of 401(k) eligiblity on financial asset holdings. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "timely-bunch",
   "metadata": {
    "papermill": {
     "duration": 0.06665,
     "end_time": "2021-03-24T14:20:21.689265",
     "exception": false,
     "start_time": "2021-03-24T14:20:21.622615",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Local Average Treatment Effects of 401(k) Participation on Net Financial Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "composite-isaac",
   "metadata": {
    "papermill": {
     "duration": 0.067774,
     "end_time": "2021-03-24T14:20:21.824102",
     "exception": false,
     "start_time": "2021-03-24T14:20:21.756328",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Interactive IV Model (IIVM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "together-editing",
   "metadata": {
    "papermill": {
     "duration": 0.066857,
     "end_time": "2021-03-24T14:20:21.957598",
     "exception": false,
     "start_time": "2021-03-24T14:20:21.890741",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now, we consider estimation of local average treatment effects (LATE) of participation with the binary instrument `e401`. As before, $Y$ denotes the outcome `net_tfa`, and $X$ is the vector of covariates.  Here the structural equation model is:\n",
    "\n",
    "\\begin{eqnarray}\n",
    "& Y = g_0(Z,X) + U, &\\quad E[U\\mid Z,X] = 0,\\\\\n",
    "& D = r_0(Z,X) + V, &\\quad E[V\\mid Z, X] = 0,\\\\\n",
    "& Z = m_0(X) + \\zeta, &\\quad E[\\zeta \\mid X] = 0.\n",
    "\\end{eqnarray}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "capable-gauge",
   "metadata": {
    "papermill": {
     "duration": 0.112727,
     "end_time": "2021-03-24T14:20:22.138131",
     "exception": false,
     "start_time": "2021-03-24T14:20:22.025404",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Constructing the data (as DoubleMLData)\n",
    "formula_flex2 = \"net_tfa ~ p401+ e401 + poly(age, 6, raw=TRUE) + poly(inc, 8, raw=TRUE) + poly(educ, 4, raw=TRUE) + poly(fsize, 2, raw=TRUE) + marr + twoearn + db + pira + hown\"\n",
    "model_flex2 = as.data.table(model.frame(formula_flex2, data))\n",
    "x_cols = colnames(model_flex2)[-c(1,2,3)]\n",
    "data_IV = DoubleMLData$new(model_flex2, y_col = \"net_tfa\", d_cols = \"p401\", z_cols =\"e401\",x_cols=x_cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "quick-fellowship",
   "metadata": {
    "papermill": {
     "duration": 21.484394,
     "end_time": "2021-03-24T14:20:43.691771",
     "exception": false,
     "start_time": "2021-03-24T14:20:22.207377",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "dml_MLIIVM = DoubleMLIIVM$new(data_IV, ml_g = lasso, \n",
    "                       ml_m = lasso_class, ml_r = lasso_class,n_folds=3, subgroups = list(always_takers = FALSE, \n",
    "                                         never_takers = TRUE))\n",
    "dml_MLIIVM$fit(store_predictions=TRUE)\n",
    "dml_MLIIVM$summary()\n",
    "lasso_MLIIVM <- dml_MLIIVM$coef\n",
    "lasso_std_MLIIVM <- dml_MLIIVM$se"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "played-rugby",
   "metadata": {
    "papermill": {
     "duration": 0.069118,
     "end_time": "2021-03-24T14:20:43.830269",
     "exception": false,
     "start_time": "2021-03-24T14:20:43.761151",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The confidence interval for the local average treatment effect of participation is given by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "healthy-lesson",
   "metadata": {
    "papermill": {
     "duration": 0.093963,
     "end_time": "2021-03-24T14:20:43.993163",
     "exception": false,
     "start_time": "2021-03-24T14:20:43.899200",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dml_MLIIVM$confint(level = 0.95)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prospective-aggregate",
   "metadata": {
    "papermill": {
     "duration": 0.07021,
     "end_time": "2021-03-24T14:20:44.133038",
     "exception": false,
     "start_time": "2021-03-24T14:20:44.062828",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Here we can also check the accuracy of the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "finished-harassment",
   "metadata": {
    "papermill": {
     "duration": 0.109791,
     "end_time": "2021-03-24T14:20:44.313881",
     "exception": false,
     "start_time": "2021-03-24T14:20:44.204090",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# variables\n",
    "y <- as.matrix(pension$net_tfa) # true observations\n",
    "d <- as.matrix(pension$p401) \n",
    "z <- as.matrix(pension$e401) \n",
    "\n",
    "# predictions\n",
    "dml_MLIIVM$params_names()\n",
    "g0_hat <- as.matrix(dml_MLIIVM$predictions$ml_g0) # predictions of g_0(z=0, X)\n",
    "g1_hat <- as.matrix(dml_MLIIVM$predictions$ml_g1) # predictions of g_0(z=1, X)\n",
    "g_hat <- z*g1_hat+(1-z)*g0_hat # predictions of g_0\n",
    "r0_hat <- as.matrix(dml_MLIIVM$predictions$ml_r0) # predictions of r_0(z=0, X)\n",
    "r1_hat <- as.matrix(dml_MLIIVM$predictions$ml_r1) # predictions of r_0(z=1, X)\n",
    "r_hat <- z*r1_hat+(1-z)*r0_hat # predictions of r_0\n",
    "m_hat <- as.matrix(dml_MLIIVM$predictions$ml_m) # predictions of m_o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "national-steering",
   "metadata": {
    "papermill": {
     "duration": 0.113011,
     "end_time": "2021-03-24T14:20:44.498082",
     "exception": false,
     "start_time": "2021-03-24T14:20:44.385071",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# cross-fitted RMSE: outcome\n",
    "lasso_y_MLIIVM <- sqrt(mean((y-g_hat)^2)) \n",
    "lasso_y_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "lasso_d_MLIIVM <- sqrt(mean((d-r_hat)^2)) \n",
    "lasso_d_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: instrument\n",
    "lasso_z_MLIIVM <- sqrt(mean((z-m_hat)^2)) \n",
    "lasso_z_MLIIVM\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "consistent-compression",
   "metadata": {
    "papermill": {
     "duration": 0.070907,
     "end_time": "2021-03-24T14:20:44.640829",
     "exception": false,
     "start_time": "2021-03-24T14:20:44.569922",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Again, we repeat the procedure for the other machine learning methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "graduate-harrison",
   "metadata": {
    "papermill": {
     "duration": 70.113632,
     "end_time": "2021-03-24T14:21:54.825289",
     "exception": false,
     "start_time": "2021-03-24T14:20:44.711657",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "### random forest ###\n",
    "\n",
    "set.seed(123)\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "dml_MLIIVM = DoubleMLIIVM$new(data_IV, ml_g = randomForest, \n",
    "                       ml_m = randomForest_class, ml_r = randomForest_class,n_folds=3, subgroups = list(always_takers = FALSE, \n",
    "                                         never_takers = TRUE))\n",
    "dml_MLIIVM$fit(store_predictions=TRUE)\n",
    "dml_MLIIVM$summary()\n",
    "forest_MLIIVM <- dml_MLIIVM$coef\n",
    "forest_std_MLIIVM <- dml_MLIIVM$se\n",
    "\n",
    "# predictions\n",
    "g0_hat <- as.matrix(dml_MLIIVM$predictions$ml_g0) # predictions of g_0(Z=0, X)\n",
    "g1_hat <- as.matrix(dml_MLIIVM$predictions$ml_g1) # predictions of g_0(Z=1, X)\n",
    "g_hat <- z*g1_hat+(1-z)*g0_hat # predictions of g_0\n",
    "r0_hat <- as.matrix(dml_MLIIVM$predictions$ml_r0) # predictions of r_0(Z=0, X)\n",
    "r1_hat <- as.matrix(dml_MLIIVM$predictions$ml_r1) # predictions of r_0(Z=1, X)\n",
    "r_hat <- z*r1_hat+(1-z)*r0_hat # predictions of r_0\n",
    "m_hat <- as.matrix(dml_MLIIVM$predictions$ml_m) # predictions of m_o\n",
    "\n",
    "# cross-fitted RMSE: outcome\n",
    "forest_y_MLIIVM <- sqrt(mean((y-g_hat)^2)) \n",
    "forest_y_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "forest_d_MLIIVM <- sqrt(mean((d-r_hat)^2)) \n",
    "forest_d_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: instrument\n",
    "forest_z_MLIIVM <- sqrt(mean((z-m_hat)^2)) \n",
    "forest_z_MLIIVM\n",
    "\n",
    "### trees ###\n",
    "\n",
    "dml_MLIIVM = DoubleMLIIVM$new(data_IV, ml_g = trees, \n",
    "                       ml_m = trees_class, ml_r = trees_class,n_folds=3, subgroups = list(always_takers = FALSE, \n",
    "                                         never_takers = TRUE))\n",
    "dml_MLIIVM$fit(store_predictions=TRUE)\n",
    "dml_MLIIVM$summary()\n",
    "tree_MLIIVM <- dml_MLIIVM$coef\n",
    "tree_std_MLIIVM <- dml_MLIIVM$se\n",
    "\n",
    "# predictions\n",
    "g0_hat <- as.matrix(dml_MLIIVM$predictions$ml_g0) # predictions of g_0(Z=0, X)\n",
    "g1_hat <- as.matrix(dml_MLIIVM$predictions$ml_g1) # predictions of g_0(Z=1, X)\n",
    "g_hat <- z*g1_hat+(1-z)*g0_hat # predictions of g_0\n",
    "r0_hat <- as.matrix(dml_MLIIVM$predictions$ml_r0) # predictions of r_0(Z=0, X)\n",
    "r1_hat <- as.matrix(dml_MLIIVM$predictions$ml_r1) # predictions of r_0(Z=1, X)\n",
    "r_hat <- z*r1_hat+(1-z)*r0_hat # predictions of r_0\n",
    "m_hat <- as.matrix(dml_MLIIVM$predictions$ml_m) # predictions of m_o\n",
    "\n",
    "# cross-fitted RMSE: outcome\n",
    "tree_y_MLIIVM <- sqrt(mean((y-g_hat)^2)) \n",
    "tree_y_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "tree_d_MLIIVM <- sqrt(mean((d-r_hat)^2)) \n",
    "tree_d_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: instrument\n",
    "tree_z_MLIIVM <- sqrt(mean((z-m_hat)^2)) \n",
    "tree_z_MLIIVM\n",
    "\n",
    "\n",
    "### boosting ###\n",
    "dml_MLIIVM = DoubleMLIIVM$new(data_IV, ml_g = boost, \n",
    "                       ml_m = boost_class, ml_r = boost_class,n_folds=3, subgroups = list(always_takers = FALSE, \n",
    "                                         never_takers = TRUE))\n",
    "dml_MLIIVM$fit(store_predictions=TRUE)\n",
    "dml_MLIIVM$summary()\n",
    "boost_MLIIVM <- dml_MLIIVM$coef\n",
    "boost_std_MLIIVM <- dml_MLIIVM$se\n",
    "\n",
    "# predictions\n",
    "g0_hat <- as.matrix(dml_MLIIVM$predictions$ml_g0) # predictions of g_0(Z=0, X)\n",
    "g1_hat <- as.matrix(dml_MLIIVM$predictions$ml_g1) # predictions of g_0(Z=1, X)\n",
    "g_hat <- z*g1_hat+(1-z)*g0_hat # predictions of g_0\n",
    "r0_hat <- as.matrix(dml_MLIIVM$predictions$ml_r0) # predictions of r_0(Z=0, X)\n",
    "r1_hat <- as.matrix(dml_MLIIVM$predictions$ml_r1) # predictions of r_0(Z=1, X)\n",
    "r_hat <- z*r1_hat+(1-z)*r0_hat # predictions of r_0\n",
    "m_hat <- as.matrix(dml_MLIIVM$predictions$ml_m) # predictions of m_o\n",
    "\n",
    "# cross-fitted RMSE: outcome\n",
    "boost_y_MLIIVM <- sqrt(mean((y-g_hat)^2)) \n",
    "boost_y_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: treatment\n",
    "boost_d_MLIIVM <- sqrt(mean((d-r_hat)^2)) \n",
    "boost_d_MLIIVM\n",
    "\n",
    "# cross-fitted RMSE: instrument\n",
    "boost_z_MLIIVM <- sqrt(mean((z-m_hat)^2)) \n",
    "boost_z_MLIIVM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "japanese-missouri",
   "metadata": {
    "papermill": {
     "duration": 0.125566,
     "end_time": "2021-03-24T14:21:55.033100",
     "exception": false,
     "start_time": "2021-03-24T14:21:54.907534",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "library(xtable)\n",
    "table <- matrix(0, 5, 4)\n",
    "table[1,1:4]   <- c(lasso_MLIIVM,forest_MLIIVM,tree_MLIIVM,boost_MLIIVM)\n",
    "table[2,1:4]   <- c(lasso_std_MLIIVM,forest_std_MLIIVM,tree_std_MLIIVM,boost_std_MLIIVM)\n",
    "table[3,1:4]   <- c(lasso_y_MLIIVM,forest_y_MLIIVM,tree_y_MLIIVM,boost_y_MLIIVM)\n",
    "table[4,1:4]   <- c(lasso_d_MLIIVM,forest_d_MLIIVM,tree_d_MLIIVM,boost_d_MLIIVM)\n",
    "table[5,1:4]   <- c(lasso_z_MLIIVM,forest_z_MLIIVM,tree_z_MLIIVM,boost_z_MLIIVM)\n",
    "rownames(table) <- c(\"Estimate\",\"Std.Error\",\"RMSE Y\",\"RMSE D\",\"RMSE Z\")\n",
    "colnames(table) <- c(\"Lasso\",\"Random Forest\",\"Trees\",\"Boosting\")\n",
    "tab<- xtable(table, digits = 2)\n",
    "tab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "animal-dancing",
   "metadata": {
    "papermill": {
     "duration": 0.078988,
     "end_time": "2021-03-24T14:21:55.191630",
     "exception": false,
     "start_time": "2021-03-24T14:21:55.112642",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We report results based on four ML methods for estimating the nuisance functions used in\n",
    "forming the orthogonal estimating equations. We find again that the estimates of the treatment effect are stable across ML methods. The estimates are highly significant, hence we would reject the hypothesis\n",
    "that the effect of 401(k) participation has no effect on financial health."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "failing-cambridge",
   "metadata": {
    "papermill": {
     "duration": 0.080777,
     "end_time": "2021-03-24T14:21:55.350975",
     "exception": false,
     "start_time": "2021-03-24T14:21:55.270198",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We might rerun the model using the best ML method for each equation to get a final estimate for the treatment effect of participation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "advisory-taxation",
   "metadata": {
    "papermill": {
     "duration": 48.36793,
     "end_time": "2021-03-24T14:22:43.798156",
     "exception": false,
     "start_time": "2021-03-24T14:21:55.430226",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "set.seed(123)\n",
    "lgr::get_logger(\"mlr3\")$set_threshold(\"warn\") \n",
    "dml_MLIIVM = DoubleMLIIVM$new(data_IV, ml_g = randomForest, \n",
    "                       ml_m = lasso_class, ml_r = lasso_class,n_folds=3, subgroups = list(always_takers = FALSE, \n",
    "                                         never_takers = TRUE))\n",
    "dml_MLIIVM$fit(store_predictions=TRUE)\n",
    "dml_MLIIVM$summary()\n",
    "best_MLIIVM <- dml_MLIIVM$coef\n",
    "best_std_MLIIVM <- dml_MLIIVM$se"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "R",
   "language": "R",
   "name": "ir"
  },
  "language_info": {
   "codemirror_mode": "r",
   "file_extension": ".r",
   "mimetype": "text/x-r-source",
   "name": "R",
   "pygments_lexer": "r",
   "version": "3.6.3"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 537.547458,
   "end_time": "2021-03-24T14:22:44.931595",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2021-03-24T14:13:47.384137",
   "version": "2.3.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
